import json
import os
from traceback import print_tb
from jax import grad
import jax.numpy as jnp
from jax import jit
import time
from matplotlib import pyplot as plt
import numpy as np
import numpy.random as npr
import jax
import jax.numpy as jnp
from jax import device_put
from jax import jit, grad, lax, random
from jax.example_libraries import optimizers
from jax.example_libraries import stax
from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus, Sigmoid, FanInSum
from jax.nn import sigmoid
from functools import partial
from jax import vmap
from flax import linen as nn
from flax.training import train_state
from flax import struct
from jax import lax

from jax import tree_util
from jax.tree_util import tree_structure
from jax.tree_util import tree_flatten, tree_unflatten

import jax.experimental.sparse as sparse

import optax
import cv2

from openTSNE import TSNE
from sklearn.manifold import Isomap

rnn_log_pth = "./data/adaptive_trajectory_optimization/rnn_state_opt_log_128_gru.npy"
trj_log_pth = "./data/adaptive_trajectory_optimization/trj_log_128_gru.json"

rnn_state_opt_log = np.load(rnn_log_pth)

print("shape of rnn_state_opt_log: ", rnn_state_opt_log.shape)

# erase the first element of rnn_state_opt_log
rnn_state_opt_log = rnn_state_opt_log[1:]

# perform t-SNE on rnn_state_opt_log
tsne_embedding = TSNE().fit(rnn_state_opt_log)
print("shape of tsne_embedding: ", tsne_embedding.shape)
# visualize the t-SNE embedding
plt.scatter(tsne_embedding[:, 0], tsne_embedding[:, 1])
plt.show()

# perform k-means on tsne_embedding to get the cluster labels
from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters=4, random_state=0).fit(tsne_embedding)
# kmeans = KMeans(n_clusters=4, random_state=0).fit(rnn_state_opt_log)
print("shape of kmeans.labels_: ", kmeans.labels_.shape)
print("kmeans.labels_: ", kmeans.labels_)
# visualize the k-means clustering result
plt.scatter(tsne_embedding[:, 0], tsne_embedding[:, 1], c=kmeans.labels_)
plt.show()

with open(trj_log_pth, "r") as f:
    trj_json = json.load(f)
    trj_log = trj_json["data"]

print("shape of rnn_state_opt_log: ", rnn_state_opt_log.shape)
print("shape of trj_log: ", len(trj_log))


c_table = np.zeros((10, 3), np.uint8)
c_table[0, :] = (255, 0, 0)
c_table[1, :] = (0, 255, 0)
c_table[2, :] = (0, 0, 255)
c_table[3, :] = (255, 255, 0)
c_table[4, :] = (255, 0, 255)
c_table[5, :] = (0, 255, 255)
c_table[6, :] = (255, 120, 53)
c_table[7, :] = (120, 255, 255)
c_table[8, :] = (255, 255, 120)
c_table[9, :] = (120, 0, 120)

# do PCA on rnn_state_opt_log
from sklearn.decomposition import PCA
pca = PCA()  
pca.fit(rnn_state_opt_log)

xt = pca.transform(rnn_state_opt_log)

# plot the first 3 dimensions of xt in 3D
fig = plt.figure()
ax = plt.axes(projection='3d')
n = len(rnn_state_opt_log)-1
k1 = 0
k2 = 1
k3 = 2
# plot the first 3 dimensions of xt in 3D by dots

ax.scatter3D(xt[:n, k1], xt[:n, k2], xt[:n, k3], c='red')
plt.show()


# # draw each element of trj_log one single opencv canvas
# img_seq = []
# task_set = []
# scale = 4
# for i in range(len(trj_log)):
#     # check if trj_log[i] is empty
#     if len(trj_log[i]) == 0:
#         img_seq.append(np.zeros((60, 60, 3), np.uint8))
#         task_set.append((0,0))
#         print("-----------------------------------------empty trj_log: ", i)
#         continue

#     img = np.zeros((60, 60, 3), np.uint8)
#     for j in range(len(trj_log[i])-1):
#         cv2.line(img, (int(trj_log[i][j][0]/scale), int(trj_log[i][j][1]/scale)), (int(trj_log[i][j+1][0]/scale), int(trj_log[i][j+1][1]/scale)), (255, 255, 255), 1)
#     # draw a bigger dot at the start point
#     # print(trj_log[i][0][0], trj_log[i][0][1])
#     # get start location
#     start_x, start_y = trj_log[i][0][0], trj_log[i][0][1]
#     # get goal location
#     goal_x, goal_y = trj_log[i][-1][0], trj_log[i][-1][1]
#     # get task
#     task = (goal_x-start_x, goal_y-start_y)
#     task_set.append(task)
#     cv2.circle(img, (int(trj_log[i][0][0]/scale), int(trj_log[i][0][1]/scale)), 2, (255, 255, 255), -1)
#     # draw a red dot at the end point
#     cv2.circle(img, (int(trj_log[i][-1][0]/scale), int(trj_log[i][-1][1]/scale)), 2, (0, 0, 255), -1)
#     # draw a white border
#     cv2.rectangle(img, (0, 0), (59, 59), (255, 255, 255), 1)
#     img_seq.append(img)

# n_img = len(img_seq)

# # rearrange img_seq on a big canvas with a grid formation
# n_row = 20
# n_col = 20
# big_img = np.zeros((n_row*60, n_col*60, 3), np.uint8)
# img_id = 0
# for i in range(n_row):
#     for j in range(n_col):
#         if img_id < n_img:
#             big_img[i*60:(i+1)*60, j*60:(j+1)*60, :] = img_seq[img_id]
#         img_id += 1

# big_img_cpy = big_img.copy()

# img_id = 0
# # draw border for each grid with color corresponding to the cluster label
# # make a color table

# num_clusters = 10

# for i in range(n_row):
#     for j in range(n_col):
#         if img_id < n_img:
#             cv2.rectangle(big_img, (j*60+3, i*60+3), ((j+1)*60-3, (i+1)*60-3), 
#             (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#             , 3)
#         img_id += 1

# # draw the big canvas
# cv2.imshow("big_img", big_img)
# cv2.waitKey(0)

# # visualize the task set on a single image
# task_img = np.zeros((60*scale*3, 60*scale*3, 3), np.uint8)
# for i in range(len(task_set)):
#     # draw a line from center to the task+center
#     center_x, center_y = task_img.shape[0]//2, task_img.shape[1]//2
#     img_id = i
#     cv2.line(task_img, (center_x, center_y), (center_x+task_set[i][0], center_y+task_set[i][1]), 
#     (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#     # (255,0,0)
#     , 1)
#     if kmeans.labels_[i] == 0:
#         print(i)
    
# cv2.imshow("task_img", task_img)
# cv2.waitKey(0)

# var_clt = []
# rnn_state_opt_log_clusters = []
# for clt in range(4):
#     # collect rnn_state_opt_log with corresponding cluster label
#     rnn_state_opt_log_cluster = []
#     img_ids = []
#     for i in range(len(kmeans.labels_)):
#         if kmeans.labels_[i] == clt:
#             rnn_state_opt_log_cluster.append(rnn_state_opt_log[i])
#             img_ids.append(i)
#     rnn_state_opt_log_clusters.append(rnn_state_opt_log_cluster)

#     # compute variance of each element of rnn_state_opt_log_cluster
#     var = np.var(rnn_state_opt_log_cluster, axis=0)
#     var_clt.append(var)
    
# # visualize the variance var_clt
# plt.figure()
# plt.plot(var_clt[0], label='cluster 0')
# plt.plot(var_clt[1], label='cluster 1')
# plt.plot(var_clt[2], label='cluster 2')
# plt.plot(var_clt[3], label='cluster 3')
# plt.legend()
# plt.show()

# var_clt_mean = np.mean(var_clt, axis=0)
# print("cross-task variance: ", var_clt_mean)
# plt.figure()
# plt.title("mean of cross-task variance")
# plt.plot(var_clt_mean)
# plt.show()

# # sort var_clt with id as lable
# var_clt = np.array(var_clt)
# var_clt = var_clt[:, np.argsort(var_clt_mean)]
# plt.figure()
# plt.plot(var_clt[0], label='cluster 0')
# plt.plot(var_clt[1], label='cluster 1')
# plt.plot(var_clt[2], label='cluster 2')
# plt.plot(var_clt[3], label='cluster 3')
# plt.legend()
# plt.show()

# var_clt_mean_sorted = np.sort(var_clt_mean)
# plt.figure()
# plt.plot(var_clt_mean_sorted)
# plt.title("sorted mean of variance of each dimension of rnn_state_opt_log")
# plt.show()

# """
# static embedding analysis
# """
# figures = []
# for v in range(5, 20):
#     var_th = v*0.01
#     print("var_th: ", var_th)
#     # collect the id of elements in var_clt_mean whose value is less than 0.1
#     var_clt_mean_less_than_0_1 = []
#     for i in range(len(var_clt_mean)):
#         if var_clt_mean[i] < var_th:
#             var_clt_mean_less_than_0_1.append(i)

#     print("var_clt_mean_less_than_0_1: ", var_clt_mean_less_than_0_1)

#     # extract the corresponding rnn_state_opt_log with id in var_clt_mean_less_than_0_1
#     rnn_state_opt_log_stationary = []
#     for i in range(len(rnn_state_opt_log)):
#         rnn_state_opt_log_stationary.append(rnn_state_opt_log[i][var_clt_mean_less_than_0_1])

#     rnn_state_opt_log_stationary = np.array(rnn_state_opt_log_stationary)
#     print(rnn_state_opt_log_stationary.shape)

#     # perform t-SNE on rnn_state_opt_log_stationary
#     tsne_embedding = TSNE().fit(rnn_state_opt_log_stationary)
#     # visualize the t-SNE embedding
#     plt.scatter(tsne_embedding[:, 0], tsne_embedding[:, 1])
#     plt.title("var_th = {}, dimensions_selected = {}".format(var_th, rnn_state_opt_log_stationary.shape[1]))
#     plt.show()

#     kmeans = KMeans(n_clusters=5, random_state=0).fit(tsne_embedding)

#     big_img = big_img_cpy.copy()
#     img_id = 0
#     for i in range(n_row):
#         for j in range(n_col):
#             if img_id < n_img:
#                 cv2.rectangle(big_img, (j*60+3, i*60+3), ((j+1)*60-3, (i+1)*60-3), 
#                 (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#                 , 3)
#             img_id += 1

#     # draw the big canvas
#     cv2.imshow("big_img", big_img)
#     cv2.waitKey(0)

#     # visualize the task set on a single image
#     task_img1 = np.zeros((60*scale*3, 60*scale*3, 3), np.uint8)
#     task_img2 = np.zeros((60*scale*3, 60*scale*3, 3), np.uint8)
#     task_img3 = np.zeros((60*scale*3, 60*scale*3, 3), np.uint8)
#     task_img4 = np.zeros((60*scale*3, 60*scale*3, 3), np.uint8)
#     img_id = 0
#     for i in range(len(task_set)):
#         # draw a line from center to the task+center
#         center_x, center_y = task_img1.shape[0]//2, task_img1.shape[1]//2
#         img_id = i
#         if kmeans.labels_[img_id] == 0:
#             cv2.line(task_img1, (center_x, center_y), (center_x+task_set[i][0], center_y+task_set[i][1]), 
#             (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#             , 1)
#         elif kmeans.labels_[img_id] == 1:
#             cv2.line(task_img2, (center_x, center_y), (center_x+task_set[i][0], center_y+task_set[i][1]), 
#             (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#             , 1)
#         elif kmeans.labels_[img_id] == 2:
#             cv2.line(task_img3, (center_x, center_y), (center_x+task_set[i][0], center_y+task_set[i][1]), 
#             (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#             , 1)
#         elif kmeans.labels_[img_id] == 3:
#             cv2.line(task_img4, (center_x, center_y), (center_x+task_set[i][0], center_y+task_set[i][1]), 
#             (int(c_table[kmeans.labels_[img_id],0]), int(c_table[kmeans.labels_[img_id],1]), int(c_table[kmeans.labels_[img_id],2]))
#             , 1)
#     # concatenate the task_img1, task_img2, task_img3, task_img4
#     task_img1 = np.concatenate((task_img1, task_img2, task_img3, task_img4), axis=1)
#     # resize the task_img1 to 0.5x
#     task_img1 = cv2.resize(task_img1, (0,0), fx=0.5, fy=0.5)

#     cv2.imshow("task_img1", task_img1)
#     cv2.waitKey(0)

